//
// Created by superuse on 2025/3/9.
//

#include "PSConsole.hpp"

#include <vector>

static void SetConsoleSize(const int width, const int height) {
    HANDLE hOut = GetStdHandle(STD_OUTPUT_HANDLE);
    if (hOut == INVALID_HANDLE_VALUE) return;

    SMALL_RECT rect;
    rect.Left = 0;
    rect.Top = 0;
    rect.Right = width - 1;
    rect.Bottom = height - 1;

    SetConsoleWindowInfo(hOut, TRUE, &rect);

    COORD coord;
    coord.X = width;
    coord.Y = height;

    SetConsoleScreenBufferSize(hOut, coord);
}

static void SetConsoleColor(const int color) {
    HANDLE hOut = GetStdHandle(STD_OUTPUT_HANDLE);
    if (hOut == INVALID_HANDLE_VALUE) return;

    SetConsoleTextAttribute(hOut, color);
}

static HANDLE hChildStd_IN_Rd = nullptr;
static HANDLE hChildStd_IN_Wr = nullptr;
static HANDLE hChildStd_OUT_Rd = nullptr;
static HANDLE hChildStd_OUT_Wr = nullptr;

static void CreatePowerShellProcess() {
    static SECURITY_ATTRIBUTES saAttr;
    saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
    saAttr.bInheritHandle = TRUE;
    saAttr.lpSecurityDescriptor = nullptr;
    if (!CreatePipe(&hChildStd_IN_Rd, &hChildStd_IN_Wr, &saAttr, 0))
        throw std::runtime_error("StdInPipe creation failed");

    if (!CreatePipe(&hChildStd_OUT_Rd, &hChildStd_OUT_Wr, &saAttr, 0))
        throw std::runtime_error("StdOutPipe creation failed");

    // 子进程
    static STARTUPINFO startup_info;
    static PROCESS_INFORMATION process_info;
    ZeroMemory(&startup_info, sizeof(startup_info));
    startup_info.cb = sizeof(startup_info); // 结构体大小
    startup_info.hStdError = hChildStd_OUT_Wr; // 错误输出重定向
    startup_info.hStdOutput = hChildStd_OUT_Wr; // 输出重定向
    startup_info.hStdInput = hChildStd_IN_Rd; // 输入重定向
    startup_info.dwFlags |= STARTF_USESTDHANDLES; // 使用标准句柄
    ZeroMemory(&process_info, sizeof(process_info));
    if (!CreateProcess(nullptr, (LPSTR) "powershell.exe",
                       nullptr,
                       nullptr,
                       TRUE, 0,
                       nullptr,
                       nullptr, &startup_info, &process_info)) {
        throw std::runtime_error("CreateProcess failed");
    }
    CloseHandle(process_info.hProcess);
    CloseHandle(process_info.hThread);
    CloseHandle(hChildStd_OUT_Wr);
    CloseHandle(hChildStd_IN_Rd);
    DWORD exitCode;
    if (GetExitCodeProcess(process_info.hProcess, &exitCode) && exitCode != STILL_ACTIVE) {
         CloseHandle(hChildStd_OUT_Wr);
        throw std::runtime_error("GetExitCodeProcess failed");
    }


}

static void ExecutePowerShellCommand(const std::string &command) {
    static DWORD dwWritten;
    WriteFile(hChildStd_IN_Wr, command.c_str(), command.length(), &dwWritten, nullptr);
    WriteFile(hChildStd_IN_Wr, "\n", 1, &dwWritten, nullptr);
    static vector<char> buffer(4096);
    static DWORD bytesRead;

    // 设置超时
    COMMTIMEOUTS timeouts = {0};
    timeouts.ReadIntervalTimeout = 50; // 50ms
    SetCommTimeouts(hChildStd_OUT_Rd, &timeouts);


    while (true) {
        if (!ReadFile(hChildStd_OUT_Rd, buffer.data() , buffer.size()-1
            , &bytesRead, nullptr)
            || bytesRead == 0) {
            break;
        }
        buffer[bytesRead] = '\0';
        cout << buffer.data() ;
    }
}

static void ExecutePS(const string &command) {
    std::string fullCommand = "powershell  -NoProfile  -Command  \"" + command + "\"";

    static STARTUPINFO si = {sizeof(si)};
    static PROCESS_INFORMATION pi;

    if (CreateProcess(
        nullptr, // 不指定模块名，使用命令行
        &fullCommand[0], // 命令行
        nullptr, // 进程句柄不可继承
        nullptr, // 线程句柄不可继承
        FALSE, // 不继承句柄
        0, // 无创建标志
        nullptr, // 使用父进程的环境块
        nullptr, // 使用父进程的当前目录
        &si, // 指向 STARTUPINFO 的指针
        &pi // 指向 PROCESS_INFORMATION 的指针
    )) {
        WaitForSingleObject(pi.hProcess, INFINITE);
        CloseHandle(pi.hProcess);
        CloseHandle(pi.hThread);
    } else {
        std::cerr << "Error: Unable to execute PowerShell command." << std::endl;
    }
}

u_int  PowershellConsole  () {
    SetConsoleSize(80, 25);
    SetConsoleTitle("Simple PowerShell Terminal");
    const HANDLE hIn = GetStdHandle(STD_INPUT_HANDLE);
    const HANDLE hOut = GetStdHandle(STD_OUTPUT_HANDLE);


    if (hIn == INVALID_HANDLE_VALUE || hOut == INVALID_HANDLE_VALUE) {
        std::cerr << "Error: Unable to get input/output handles." << std::endl;
        return 1;
    }

    static DWORD prevMode;
    GetConsoleMode(hIn, &prevMode);
    SetConsoleMode(hIn, ENABLE_PROCESSED_INPUT | ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);
    SetConsoleColor(FOREGROUND_GREEN | FOREGROUND_INTENSITY);
    try {
        CreatePowerShellProcess();
        static string input;
        static COMMTIMEOUTS timeouts = {0};
        timeouts.ReadIntervalTimeout = 50; // 50ms
        SetCommTimeouts(hChildStd_OUT_Rd, &timeouts);
        while (true) {
            cout << "$PS> ";
            getline(cin, input);
            if (input == "exit" || input == "quit" || input == "q") { break; }
            ExecutePowerShellCommand(input);
        }


    } catch (const std::exception &e) {
        cerr << "Terminal Error: " << e.what() << endl;
    }

    SetConsoleMode(hIn, prevMode);
    CloseHandle(hChildStd_IN_Wr);
    CloseHandle(hChildStd_OUT_Rd);
    return 0;
}

